package main

import (
	"bytes"
	"embed"
	"fmt"
	"image"
	"io"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/otiai10/gosseract/v2"
	"gocv.io/x/gocv"
)

//go:embed frontend-dist/*
var staticFiles embed.FS

type OCRJob struct {
	Data []byte
	Done chan string
	Err  chan error
}

type OCRWorker struct {
	client *gosseract.Client
}

func (w *OCRWorker) process(job OCRJob) {
	mat, err := preprocess(job.Data)
	if err != nil {
		job.Err <- err
		return
	}
	defer mat.Close()

	buf, err := gocv.IMEncode(".png", mat)
	if err != nil {
		job.Err <- err
		return
	}
	defer buf.Close()

	if err := w.client.SetImageFromBytes(buf.GetBytes()); err != nil {
		job.Err <- err
		return
	}
	text, err := w.client.Text()
	if err != nil {
		job.Err <- err
		return
	}
	job.Done <- text
}

func preprocess(imgBytes []byte) (gocv.Mat, error) {
	mat, err := gocv.IMDecode(imgBytes, gocv.IMReadColor)
	if err != nil || mat.Empty() {
		return gocv.Mat{}, fmt.Errorf("invalid image: %v", err)
	}
	gray := gocv.NewMat()
	gocv.CvtColor(mat, &gray, gocv.ColorBGRToGray)
	blur := gocv.NewMat()
	gocv.GaussianBlur(gray, &blur, image.Pt(3, 3), 0, 0, gocv.BorderDefault)
	bin := gocv.NewMat()
	gocv.AdaptiveThreshold(blur, &bin, 255, gocv.AdaptiveThresholdMean, gocv.ThresholdBinary, 15, 10)
	return bin, nil
}

func startWorkerPool(numWorkers int, jobs <-chan OCRJob) {
	for i := 0; i < numWorkers; i++ {
		go func(id int) {
			client := gosseract.NewClient()
			defer client.Close()
			client.SetLanguage("chi_sim", "eng")
			client.SetPageSegMode(gosseract.PSM_AUTO)
			worker := OCRWorker{client: client}
			for job := range jobs {
				worker.process(job)
			}
		}(i)
	}
}

func main() {
	numWorkers := runtime.NumCPU()
	fmt.Printf("启动 OCR 服务，worker 数量：%d\n", numWorkers)

	jobs := make(chan OCRJob, 32)
	startWorkerPool(numWorkers, jobs)

	r := gin.Default()

	// API 路由
	r.POST("/ocr", func(c *gin.Context) {
		file, err := c.FormFile("image")
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "missing image file"})
			return
		}
		f, _ := file.Open()
		defer f.Close()
		imgBytes, _ := io.ReadAll(f)

		job := OCRJob{Data: imgBytes, Done: make(chan string, 1), Err: make(chan error, 1)}
		jobs <- job

		select {
		case text := <-job.Done:
			c.JSON(http.StatusOK, gin.H{"text": text})
		case err := <-job.Err:
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
	})

	// 暴露 swagger.yaml （作为 /swagger.json 也行）
	r.GET("/swagger.yaml", func(c *gin.Context) {
		data, err := staticFiles.ReadFile("swagger.yaml")
		if err != nil {
			c.String(http.StatusNotFound, "swagger not found")
			return
		}
		c.Data(http.StatusOK, "application/yaml; charset=utf-8", data)
	})

	// 静态文件服务
	r.NoRoute(func(c *gin.Context) {
		path := c.Request.URL.Path
		if path == "/" {
			path = "/index.html"
		}
		data, err := staticFiles.ReadFile("frontend-dist" + path)
		if err != nil {
			c.Status(404)
			return
		}
		http.ServeContent(c.Writer, c.Request, path, time.Time{}, bytes.NewReader(data))
	})

	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}
	r.Run(":" + port)
}
